#include "DTI.h"

#include <teem/ten.h>


CDTI::CDTI(void)
{
	m_pDTI = NULL;
	m_w = m_h = m_d = m_n = 0;
	m_order = 0;
}


CDTI::~CDTI(void)
{
	SafeDeleteArray(m_pDTI);
	m_w = m_h = m_d = m_n = 0;
	m_order = 0;
}

int CDTI::OpenDTIFile(const char *pathname, FILE *fp)
{
	if (pathname == NULL)
		return -1;

	fprintf(fp, "Reading DTI file %s ...", pathname);

	Nrrd *dti = nrrdNew();
	if (nrrdLoad(dti, pathname, NULL)) {
		fprintf(fp, "cannot load file %s:\n%s\n", pathname, biffGetDone(NRRD));
		return -2;			// cannot load file
	}

	m_n = (uint)(dti->axis[0].size);
	m_w = (uint)(dti->axis[1].size);
	m_h = (uint)(dti->axis[2].size);
	m_d = (uint)(dti->axis[3].size);

	if (m_n != 7) {
		fprintf(fp, "Not a DTI file!\n");
		return -3;
	}

	m_order = 2;

	m_pDTI = new float[m_w*m_h*m_d*m_n];
	memcpy(m_pDTI, dti->data, sizeof(float)*m_w*m_h*m_d*m_n);
	nrrdNuke(dti);

	fprintf(fp, "Success.\n");
	fprintf(fp, "Width: %d, Height: %d, Depth: %d, Order: %d\n", 
		m_w, m_h, m_d, m_order);
	return 0;
}

int CDTI::SaveDTIFile(const char *pathname, FILE *fp)
{
	if (pathname == NULL || m_pDTI == NULL)
		return -1;

	Nrrd *dst = nrrdNew();
	size_t ss[4];
	ss[0] = m_n;
	ss[1] = m_w;
	ss[2] = m_h;
	ss[3] = m_d;

	nrrdAlloc_nva(dst, nrrdTypeFloat, 4, ss);
	float *data = (float *)(dst->data);
	uint size = m_w * m_h * m_d * m_n;
	memcpy(data, m_pDTI, sizeof(float)*size);

	nrrdAxisInfoSet_va(dst, nrrdAxisInfoSpacing, 1.0, 1.0, 1.0, 1.0);
	nrrdAxisInfoSet_va(dst, nrrdAxisInfoKind, nrrdKindList, nrrdKindSpace, nrrdKindSpace, nrrdKindSpace);
	nrrdAxisInfoSet_va(dst, nrrdAxisInfoUnits, "", "mm", "mm", "mm");

	if (nrrdSave(pathname, dst, NULL)) {
		fprintf(fp, "cannot save file %s:\n%s\n", pathname, biffGetDone(NRRD));
		return -1;
	}

	nrrdNuke(dst);

	return 0;
}

int CDTI::CreateDTI(const uint w, const uint h, const uint d, const float *dti)
{
	if (m_pDTI != NULL)
		return -1;

	uint size = w * h * d * 7;
	if (size == 0)
		return -2;

	m_pDTI = new float[size];
	memset(m_pDTI, 0, sizeof(float)*size);
	m_w = w;
	m_h = h;
	m_d = d;
	m_n = 7;

	if (dti != NULL) {
		memcpy(m_pDTI, dti, sizeof(float)*size);
	}

	return 0;
}


void CDTI::SetDTI(const uint x, const uint y, const uint z, float *dti)
{
	uint index = ((z * m_h + y) * m_w + x) * m_n;
	for (uint i = 0; i < m_n; ++i) {
		m_pDTI[index+i] = dti[i];
	}
}

void CDTI::GetDTIFast(const uint x, const uint y, const uint z, float *dti)
{
	uint index = ((z * m_h + y) * m_w + x) * m_n;
	for (uint i = 0; i < m_n; ++i) {
		dti[i] = m_pDTI[index+i];
	}
}
void CDTI::GetDTIFast(const uint index, float *dti)
{
	for (uint i = 0; i < m_n; ++i) {
		dti[i] = m_pDTI[index*m_n+i];
	}
}


void CDTI::GetDTI(const uint x, const uint y, const uint z, float *dti)
{
	memset(dti, 0, sizeof(float)*m_n);

	if (x < 0 || x > (m_w-1) ||
		y < 0 || y > (m_h-1) ||
		z < 0 || z > (m_d-1)) {
			return;
	}

	uint index = ((z * m_h + y) * m_w + x) * 7;

	for (uint i = 0; i < m_n; ++i) {
		dti[i] = m_pDTI[index+i];
	}
}

void CDTI::GetDTI(const float x, const float y, const float z, float *dti)
{
	memset(dti, 0, sizeof(float)*m_n);

	if (x < 0.0f || x > (float)(m_w-1) ||
		y < 0.0f || y > (float)(m_h-1) ||
		z < 0.0f || z > (float)(m_d-1)) {
			return;
	}

	uint x1 = (uint)floor(x);
	uint y1 = (uint)floor(y);
	uint z1 = (uint)floor(z);

	uint x2 = (uint)ceil(x);
	uint y2 = (uint)ceil(y);
	uint z2 = (uint)ceil(z);

	float a = x - (float)x1;
	float b = y - (float)y1;
	float c = z - (float)z1;

	uint i1 = ((z1 * m_h + y1) * m_w + x1)*7;
	uint i2 = ((z1 * m_h + y1) * m_w + x2)*7;
	uint i3 = ((z1 * m_h + y2) * m_w + x2)*7;
	uint i4 = ((z1 * m_h + y2) * m_w + x1)*7;

	uint i5 = ((z2 * m_h + y1) * m_w + x1)*7;
	uint i6 = ((z2 * m_h + y1) * m_w + x2)*7;
	uint i7 = ((z2 * m_h + y2) * m_w + x2)*7;
	uint i8 = ((z2 * m_h + y2) * m_w + x1)*7;

	for (uint i = 0; i < m_n; ++i) {
		float v1 = m_pDTI[i1+i];
		float v2 = m_pDTI[i2+i];
		float v3 = m_pDTI[i3+i];
		float v4 = m_pDTI[i4+i];
		float v5 = m_pDTI[i5+i];
		float v6 = m_pDTI[i6+i];
		float v7 = m_pDTI[i7+i];
		float v8 = m_pDTI[i8+i];

		float temp = v1 + a * (v2 - v1) + b * (v4 - v1) + c * (v5 - v1) + 
			a * b * (v1 - v2 + v3 - v4) + 
			a * c * (v1 - v2 + v6 - v5) + 
			b * c * (v1 - v4 + v8 - v5) + 
			a * b * c * (-v1 + v2 - v3 + v4 + v5- v6 + v7 - v8);
		dti[i] = temp;
	}
}

void CDTI::GetDTI(const CPoint3F &p, float *dti)
{
	return GetDTI(p.m_x, p.m_y, p.m_z, dti);
}

void CDTI::GetDTI(const CPoint3F *p, float *dti)
{
	return GetDTI(p->m_x, p->m_y, p->m_z, dti);
}

void CDTI::DTI2Direction(const float *dti, CVector3F *dir) const
{
	/* using teem to compute the eigenvector from the 2nd order tensor */
	double m[9], eigenValues[3], eigenVectors[9], trc, iso[9];
	int ret;

	TEN_T2M(m, dti);
	trc = ELL_3M_TRACE(m)/3.0;
	ELL_3M_SCALE_SET(iso, -trc, -trc, -trc);
	ELL_3M_ADD2(m, m, iso);

	ret = ell_3m_eigensolve_d(eigenValues, eigenVectors, m, AIR_TRUE);

	if (ell_cubic_root_single_double == ret) {
		/* this was added to fix a stupid problem with very nearly
		isotropic glyphs, used for demonstration figures */
		if (eigenValues[0] == eigenValues[1]) {
			ELL_3V_CROSS(eigenVectors+6, eigenVectors+0, eigenVectors+3);
		} else {
			ELL_3V_CROSS(eigenVectors+0, eigenVectors+3, eigenVectors+6);
		}
	}

	dir->m_x = (float)(eigenVectors[0]);
	dir->m_y = (float)(eigenVectors[1]);
	dir->m_z = (float)(eigenVectors[2]);
}

float CDTI::DTI2CL(const float *dti) const
{
	/* using teem to compute the eigenvector from the 2nd order tensor */
	double m[9], eigenValues[3], trc, iso[9];
	int ret;

	TEN_T2M(m, dti);
	trc = ELL_3M_TRACE(m)/3.0;
	ELL_3M_SCALE_SET(iso, -trc, -trc, -trc);
	ELL_3M_ADD2(m, m, iso);

    ret = ell_3m_eigenvalues_d(eigenValues, m, AIR_TRUE);
    ELL_3V_SET(eigenValues, eigenValues[0] + trc, eigenValues[1] + trc, eigenValues[2] + trc);

	for (int i = 0; i < 3; ++i) {
		if (eigenValues[i] < 0.0f)
			eigenValues[i] = 0.0f;
	}
	float cl = 0.0f;
	if ((eigenValues[0] + eigenValues[1] + eigenValues[2]) > 0.0f)
		cl = (eigenValues[0] - eigenValues[1]) / (eigenValues[0] + eigenValues[1] + eigenValues[2]);
	else
		cl = 0.0f;

	if (cl > 1.0f)
		cl = 1.0f;
	if (cl < 0.0f)
		cl = 0.0f;

	return cl;
}

float CDTI::DTI2MD(const float *dti) const
{
	/* using teem to compute the eigenvector from the 2nd order tensor */
	double m[9], eigenValues[3], trc, iso[9];
	int ret;

	TEN_T2M(m, dti);
	trc = ELL_3M_TRACE(m)/3.0;
	ELL_3M_SCALE_SET(iso, -trc, -trc, -trc);
	ELL_3M_ADD2(m, m, iso);

    ret = ell_3m_eigenvalues_d(eigenValues, m, AIR_TRUE);
    ELL_3V_SET(eigenValues, eigenValues[0] + trc, eigenValues[1] + trc, eigenValues[2] + trc);

	for (int i = 0; i < 3; ++i) {
		if (eigenValues[i] < 0.0f)
			eigenValues[i] = 0.0f;
	}
	float md = 0.0f;
	md = (eigenValues[0] + eigenValues[1] + eigenValues[2]) / 3.0f;
	return md;
}

float CDTI::DTI2CP(const float *dti) const
{
	/* using teem to compute the eigenvector from the 2nd order tensor */
	double m[9], eigenValues[3], trc, iso[9];
	int ret;

	TEN_T2M(m, dti);
	trc = ELL_3M_TRACE(m)/3.0;
	ELL_3M_SCALE_SET(iso, -trc, -trc, -trc);
	ELL_3M_ADD2(m, m, iso);

    ret = ell_3m_eigenvalues_d(eigenValues, m, AIR_TRUE);
    ELL_3V_SET(eigenValues, eigenValues[0] + trc, eigenValues[1] + trc, eigenValues[2] + trc);

	for (int i = 0; i < 3; ++i) {
		if (eigenValues[i] < 0.0f)
			eigenValues[i] = 0.0f;
	}
	float cp = 0.0f;
	if ((eigenValues[0] + eigenValues[1] + eigenValues[2]) > 0.0f)
		cp = 2.0f * (eigenValues[1] - eigenValues[2]) / (eigenValues[0] + eigenValues[1] + eigenValues[2]);
	else
		cp = 0.0f;
	return cp;
}

//float CDTI::DTI2MD(const float *dti) const
//{
//	float t[6];
//	t[0] = dti[1]; t[1] = dti[2]; t[2] = dti[3]; t[3] = dti[4]; t[4] = dti[5]; t[5] = dti[6];
//
//	float md = t[0] + t[3] + t[5];
//	return md;
//}

float CDTI::DTI2FA(const float *dti) const
{
	float t[6];
	t[0] = dti[1]; t[1] = dti[2]; t[2] = dti[3]; t[3] = dti[4]; t[4] = dti[5]; t[5] = dti[6];

	float mean = (t[0] + t[1] + t[2] + t[3] + t[4] + t[5]) / 6.0f;
	if (mean == 0.0f)
		return 0.0f;

	t[0] = t[0] / mean;
	t[1] = t[1] / mean;
	t[2] = t[2] / mean;
	t[3] = t[3] / mean;
	t[4] = t[4] / mean;
	t[5] = t[5] / mean;

	float cross = t[1] * t[1] + t[2] * t[2] + t[4] * t[4];
	float j2 = t[0] * t[3] + t[3] * t[5] + t[5] * t[0] - cross;
	float j4 = t[0] * t[0] + t[3] * t[3] + t[5] * t[5] + 2.0f * cross;
	float fa = sqrtf((j4-j2)/j4);
	return fa;
}


void CDTI::DTI2Eigen(const float *dti, float *eval, float *evec) const
{
	/* using teem to compute the eigenvector from the 2nd order tensor */
	double m[9], eigenValues[3], eigenVectors[9], trc, iso[9];
	int ret;

	TEN_T2M(m, dti);
	trc = ELL_3M_TRACE(m)/3.0;
	ELL_3M_SCALE_SET(iso, -trc, -trc, -trc);
	ELL_3M_ADD2(m, m, iso);

	ret = ell_3m_eigensolve_d(eigenValues, eigenVectors, m, AIR_TRUE);
	ELL_3V_SET(eval, eigenValues[0] + trc, eigenValues[1] + trc, eigenValues[2] + trc);

	if (ell_cubic_root_single_double == ret) {
		/* this was added to fix a stupid problem with very nearly
		isotropic glyphs, used for demonstration figures */
		if (eigenValues[0] == eigenValues[1]) {
			ELL_3V_CROSS(eigenVectors+6, eigenVectors+0, eigenVectors+3);
		} else {
			ELL_3V_CROSS(eigenVectors+0, eigenVectors+3, eigenVectors+6);
		}
	}

	for (int i = 0; i < 3; ++i) {
		if (eval[i] < 0.0f) 
			eval[i] = 0.0f;
	}

	for (int i = 0; i < 9; ++i) {
		evec[i] = eigenVectors[i];
	}
}